"""
简单移动平均
"""
import pandas as pd
import matplotlib.pyplot as plt


def simple_ma(series, window):
    """
    简单移动平均
    :param series:
    :param window: 移动窗口
    :return:
    """
    series_mean = series.rolling(window=window).mean()
    # series_mean.columns = []
    return series_mean.tolist()


def sma(series):
    """
    简单移动平均的参数比较
    :param series:
    :return:
    """
    if not isinstance(series, pd.Series):
        series = pd.Series(series)

    rs_arr = []
    windows = [3, 7, 10]
    for window in windows:
        mid_series_mean = simple_ma(series=series, window=window)
        rs_arr.append(mid_series_mean)
    df_rs = pd.DataFrame(rs_arr).T
    df_rs.columns = ['ma_{}'.format(num) for num in windows]

    df_rs['raw'] = pd.DataFrame(series.tolist(), index=df_rs.index)

    df_rs.plot()
    plt.show()
